#from sqlalchemy.inspection import _self_inspects

__author__ = 'xiaosu'
# -*- coding: utf-8 -*-

import pexpect
import logging
import re
from multiprocessing import Pool,cpu_count,Manager
import os
import socket,time,threading
from octopus import cache
basedir = os.path.abspath(os.path.dirname(__file__))




class host_command(object):


    def __init__(self, command_token,count_dict,hostname,username,password,command,audit,enabled,serverport='22',rootpassword="",run_as_root=False,timeout='60'):

        self.hostname = hostname
        self.username = username
        self.password = password
        self.audit = audit
        self.enabled = enabled
        self.serverport = serverport
        self.rootpassword = rootpassword
        self.command = command
        self.run_as_root=run_as_root
        self.timeout=timeout
        self.command_token=command_token
        self.count_dict=count_dict
        self.logger=logging.getLogger("octopus")
        self.logger.setLevel(logging.INFO)
        self.fh=logging.FileHandler(os.path.join(basedir,'./log/run_command_log/')+self.hostname+'-'+self.username+'.log')
        self.fh.setFormatter(logging.Formatter('%(asctime)s    %(levelname)s  %(message)s'))
        self.fh.setLevel(logging.INFO)
        self.logger.addHandler(self.fh)




    def run(self):

        if self.enabled == False:
            return self.hostname+" is disabled , Please enable it!"
        ssh_newkey = 'Are you sure you want to continue connecting'
        self._child = pexpect.spawn('ssh -o "NumberOfPasswordPrompts=1"  -o "RSAAuthentication=no" -o "PubkeyAuthentication=no" -p %s -l %s %s '%(self.serverport,self.username, self.hostname),timeout=self.timeout)
        self.logger.info('ssh -o "NumberOfPasswordPrompts=1"  -o "RSAAuthentication=no" -o "PubkeyAuthentication=no" -p %s -l %s %s '%(self.serverport,self.username, self.hostname))


        i = self._child.expect([pexpect.TIMEOUT, 'Connection refused', ssh_newkey, 'password: ',pexpect.EOF])
        if i == 0 or i == 4 :
            # print 'ERROR!'
            # print 'SSH could not login. Here is what SSH said:'
            # print self._child.before, self._child.after
            self.logger.error("pexpect timeout, maybe no route to host "+self.hostname)

            return "no route to host \n"

        if i == 1:
            # print 'Host is unreachable,maybe server is down '
            # print 'SSH could not login. Here is what SSH said:'
            # print self._child.before, self._child.after
            self.logger.error("Connection refused, maybe "+self.hostname+"'s ssh port is wrong !")

            return  "Connection refused, maybe "+self.hostname+"'s ssh port is wrong ! \n"

        if i == 2: # SSH does not have the public key. Just accept it.
            self._child.sendline ('yes')
            #print self._child.before, self._child.after
            self.logger.info("add server's public key")
            i = self._child.expect([pexpect.TIMEOUT, 'password: '])
            if i == 0: # Timeout
                # print 'ERROR!'
                # print 'SSH could not login. Here is what SSH said:'
                # print self._child.before, self._child.after
                self.logger.error("pexpect timeout,server doesn't response after adding the public key")

                return "server doesn't response after adding the public key \n"

        self._child.sendline(self.password)
        p=self._child.expect([pexpect.TIMEOUT,'Permission denied', '\$'])
#        p=self._child.expect([pexpect.TIMEOUT,'Permission denied', '\$','\#'])
        if p == 0:
            #print 'password wrong'
            self.logger.error("timeout after sending password")
            return "timeout after sending password \n"
        if p == 1:
            self.logger.error(self.username+"' password is wrong")
            return  self.username+"' password is wrong \n"
        if p == 2:
            self.logger.info(self.hostname+" "+self.username+"'s password is OK")
            self.logger.info("as "+self.username+" run command:  "+self.command )
            if  self.run_as_root == False :
                self._child.sendline('export PS1="[run_with_octopus]\\\$"')       #make expect easier
                self._child.expect(["\[run_with_octopus\]\$",pexpect.TIMEOUT])
                self._child.sendline("alias rm='rm' && alias mv='mv' && alias cp='cp' ")  #avoid  interact
                self._child.expect(["\[run_with_octopus\]\$",pexpect.TIMEOUT])
                self._child.sendline(self.command)
                t=self._child.expect(["\[run_with_octopus\]\$",pexpect.TIMEOUT])
                if t == 0:
                    self._child.sendline("exit")
                if t == 1:
                    return "command excution timeout !"
                result=self._child.before
                tmp=result.split("\n")
                tmp.pop()
                del tmp[0]
                # tmp[0]=tmp[0].lstrip()
                # tmp.insert(0,"======> The command is :")
                # tmp.insert(2,"======> The result is :")
                result="\n".join(tmp)
                self.logger.info("result: "+result)
                return result
            elif self.rootpassword == "":
                return "you haven't added root's password in octopus!"
            else:
                self._child.sendline("export LANG=c")
                self._child.sendline("su - root")
                self._child.expect("assword:")
                self._child.sendline(self.rootpassword)
                r=self._child.expect(["\#","Authentication failure","incorrect password"])
                if r == 0 :
                    self.logger.info(self.hostname+" root's password is OK")
                    self.logger.info("as root run command:  "+self.command )
                    self._child.sendline('export PS1="[run_with_octopus]\\\$"')    #make expect easier
                    self._child.expect(["\[run_with_octopus\]\#",pexpect.TIMEOUT])
                    self._child.sendline("alias rm='rm' && alias mv='mv' && alias cp='cp' ")  #avoid  interact
                    self._child.expect(["\[run_with_octopus\]\#",pexpect.TIMEOUT])
                    self._child.sendline(self.command)
                    t=self._child.expect(["\[run_with_octopus\]\#",pexpect.TIMEOUT])
                    if t == 0:
                        self._child.sendline("exit")
                    if t == 1:
                        return "command excution timeout !"
                    result= self._child.before
                    tmp=result.split("\n")
                    tmp.pop()
                    del tmp[0]
                    result="\n".join(tmp)
                    self.logger.info("result: \n"+result)
                    self._child.sendline("exit")
                    return result
                if r == 1 or r ==2:
                    self.logger.info(self.hostname+"root's password is incorrect")
                    self._child.sendline("exit")
                    return  self.hostname+" root's password is incorrect \n"

    def __del__(self):
        self.logger.removeHandler(self.fh)
        self.count_dict[self.command_token]+=1



def ip_to_list(ipaddress):
    """
    parse string to IP address Pool
    """
    ip_pool=[]
    ip= re.compile("^(\d{1,2}|1\d\d|2[0-4]\d|25[0-5])(\.(\d{1,2}|1\d\d|2[0-4]\d|25[0-5])){2}\.")
    ip_group = ip.sub("",ipaddress)
    tmp = ip.match(ipaddress)
    prefix=tmp.group(0)
    single_check=re.split('[,-]',ip_group)
    single_re=re.compile("^(\d{1,2}|1\d\d|2[0-4]\d|25[0-5])$")
    for i in single_check:
        if single_re.match(i) == None:
            return "invalid IP address"
    iplist=ip_group.split(",")
    # print iplist
    for i in iplist:
        try:
            n=i.split("-")
            first=int(n[0])
            last=int(n[1])
            for ip in xrange(first,last+1):
                ip_pool.append(prefix+str(ip))
        except Exception as e:
            if str(e) == "list index out of range":
                ip_pool.append(prefix+str(i))
    return ip_pool




class multi_run(object):


    def __init__(self,host_list,command_token):

        self.host_list=host_list
        self.command_token=command_token
        self.finished_count="finished_count_"+command_token
        self.result=[]
        self.new_pool=[]
        for i in host_list:
            self.new_pool.append([i,self])

    def run_one_host(self,hostinfo):

        self.result.append(hostinfo[0])
        myhost=host_command(self.command_token,self.count_dict,hostinfo[0],hostinfo[1],hostinfo[2],hostinfo[3],hostinfo[4],hostinfo[5],hostinfo[6],hostinfo[7],hostinfo[8],hostinfo[9])
        self.result.append(myhost.run())

    def get_finished_count(self):

        finished_count_tmp=0
        while finished_count_tmp != len(self.host_list):
            interval=0.2
            finished_count=self.count_dict[self.command_token]
            try:
                cache.set(self.finished_count,finished_count)
            except Exception as e:
                pass
            time.sleep(interval)
            finished_count_tmp=cache.get(self.finished_count)

    def para_run(self):

        manager=Manager()
        self.count_dict=manager.dict()
        self.count_dict[self.command_token]=0
        pool=Pool(cpu_count())
        #start get_finished_count
        t=threading.Thread(target=self.get_finished_count,args=())
        t.start()
        #start ok
        pre_result=pool.map(unwrap_self_run_one_host,self.new_pool)
        pool.close()
        pool.join()
        tmp1=[]
        for i in pre_result:
            if i not in tmp1:
                tmp1.append(i)
        tmp2=[]
        ip_all=[]
        result_dict={}
        final_result=[]
        for i in tmp1:
            count= len(i)/2
            for x in xrange(count):
                tmp2.append([i[x*2+0],i[x*2+1]])
        for i in tmp2:
            ip_all.append(i[0])
        for i in tmp2:
            result_dict[i[0]]=i
        for i in sorted(ip_all,key=socket.inet_aton):
            final_result.append(result_dict[i])
        return final_result

    def __del__(self):
        self.count_dict.pop(self.command_token)

def unwrap_self_run_one_host(mylist):
    """
    pool.map doesn't support class.function,thus ,add a outer function for multi_add.add_one_host
    """
    group=mylist[1]  #group is an object
    hostinfo=mylist[0]
    group.run_one_host(hostinfo)
    return group.result




if __name__ == '__main__':
    test_host=host_command('192.168.3.105','username','password','whoami',audit=True,enabled=True,rootpassword="redhat",run_as_root=False)
    print test_host.run()